/// Test global to global memory copy bandwidth.

#include "benchmark_config.h"
#include "cudaDeviceQuery.h"
#include "gpu_timer.h"
#include <cuda.h>
#include <stdio.h>
#include <stdlib.h>
#include <vector>

#define CEIL_DIV(x, y) (((x) + (y)-1) / (y))
typedef float MyFloat4 __attribute__((ext_vector_type(4), aligned(16)));
#define FLOAT4(pointer) (reinterpret_cast<MyFloat4 *>(&(pointer))[0])

typedef void (*CpyKernel)(uint8_t *, uint8_t *, size_t);

// Shape: (iter_max, gridDim.x, InstPerThread, blockDim.x, 16)
// Index: (iter,     blockIdx.x,           i  ,threadIdx.x, 1 )
template <size_t InstPerThread>
__global__ void cpy_kernel_align16_rwk(uint8_t *dst, uint8_t *src, size_t n) {
  size_t stride_inst = 16 * blockDim.x;
  size_t stride_block = stride_inst * InstPerThread;
  size_t stride_iter = stride_block * gridDim.x;
  size_t offset = blockIdx.x * stride_block + threadIdx.x * 16;
  size_t iter_max =
      CEIL_DIV(n, blockDim.x * gridDim.x * InstPerThread * 128 / 8);
  MyFloat4 temp[InstPerThread];
  for (size_t iter = 0; iter < iter_max - 1; iter++) {
#pragma unroll
    for (int i = 0; i < InstPerThread; i++) {
      temp[i] = FLOAT4(src[offset + i * stride_inst]);
    }
#pragma unroll
    for (int i = 0; i < InstPerThread; i++) {
      FLOAT4(dst[offset + i * stride_inst]) = temp[i];
    }
    offset += stride_iter;
  }
  // Process tail of data
  while (offset < n) {
    FLOAT4(dst[offset]) = FLOAT4(src[offset]);
    offset += stride_block;
  }
}

float testKernel(CpyKernel kernel, uint8_t *dst, uint8_t *src, size_t n,
                 int num_blocks) {
  const int num_thread = 8 * kWarpSize;
  // Warmup
  for (int i = 0; i < WARMUPNUM; ++i)
    kernel<<<num_blocks, num_thread>>>(dst, src, n);
  cudaDeviceSynchronize();
  checkGPULastErrors();

  GpuTimer timer;
  float total_time = 0.f;
  for (int i = 0; i < 20; i++) {
    timer.start();
    kernel<<<num_blocks, num_thread>>>(dst, src, n);
    timer.stop_and_wait();
    total_time += timer.duration();
  }

  return total_time / 20;
}

int main(int argc, char **argv) {
  initializeDeviceProp(0);
  uint8_t *a_d{nullptr}, *b_d{nullptr};
  size_t num_data{gDeviceParams.NumSM * 1024 * 1024 * 10};
#if gAlign
  size_t size_data{(num_data + 512)};
#else
  size_t size_data = num_data;
#endif

  cudaMalloc((void **)&a_d, size_data);
  cudaMalloc((void **)&b_d, size_data);
  cudaMemset(a_d, 0xff, size_data);
  int num_blocks = 208;
  constexpr int kInstPerThread = 16;
  float bw;
  float kernel_time = 0.f;
  kernel_time = testKernel(cpy_kernel_align16_rwk<kInstPerThread>, b_d, a_d,
                           size_data, num_blocks);
  bw = ((float)size_data / kernel_time / 1024 / 1024);
  printf("Average kernel BW: %f GB/s\n", bw);

  cudaFree(a_d);
  cudaFree(b_d);
  return 0;
}
